package dnet

import (
	"dgo/goutils/utils"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

type ProxyTcpRec struct {
	frontend_sendsize int64
	backend_recvsize  int64
	disFlag           int32
	StartT            time.Time
	FrontendAddr      string
	BackendAddr       string
	Frontend_conn     net.Conn
	Backend_conn      net.Conn

	Tag       interface{}
	DisReason string
	sessionid string
}

func (this *ProxyTcpRec) GetSessionId() string {
	return this.sessionid
}

func (this *ProxyTcpRec) GetBackendInfo() string {
	return fmt.Sprintf("%s,%s", this.BackendAddr, utils.HumanFilesize(this.backend_recvsize))
}

func (this *ProxyTcpRec) GetFrontendInfo() string {
	return fmt.Sprintf("%s,%s", this.FrontendAddr, utils.HumanFilesize(this.frontend_sendsize))
}

func (this *ProxyTcpRec) GetFrontendSendSize() int64 {
	return this.frontend_sendsize
}

func (this *ProxyTcpRec) RequestDisconnect(reason string) {
	if atomic.CompareAndSwapInt32(&this.disFlag, 0, 1) {
		this.DisReason = reason
	}
	if this.Frontend_conn != nil {
		this.Frontend_conn.Close()
	}

	if this.Backend_conn != nil {
		this.Backend_conn.Close()
	}
}

/*
*

	代理
*/
type ProxyTcp struct {
	OnlineCnt     int32
	terminated    bool
	closeW        *sync.WaitGroup
	frontend      net.Listener
	frontend_addr string
	backend_addr  string
	Sessions      utils.GSyncMapGroup
	OnErrFunc     func(err error)
	OnListenSucc  func()
	/**
	阻塞调用:
	返回值:
	 0: 表示没有处理, 继续执行 (默认的连接处理)
	 1: 表示已经处理, 不要继续执行 (需要自己调用: OnProxyConnCb)
	*/
	OnHandleFrontendConnFunc func(frontendconn net.Conn) int

	/**
	  拨号Dial
	*/
	OnDialBackend            func(sender *ProxyTcpRec) (net.Conn, error)
	OnFrontendRecv           func(sender *ProxyTcpRec, buf []byte)
	OnFrontendRecvBufConvert func(sender *ProxyTcpRec, buf []byte) []byte
	OnBackendRecv            func(sender *ProxyTcpRec, buf []byte)
	OnBackendRecvBufConvert  func(sender *ProxyTcpRec, buf []byte) []byte
	OnDisconnected           func(sender *ProxyTcpRec)
}

func NewProxyTcp() *ProxyTcp {
	rval := &ProxyTcp{
		closeW:     new(sync.WaitGroup),
		terminated: true,
	}
	return rval
}

func (this *ProxyTcp) CheckFrontend() net.Listener {
	if this.frontend == nil {
		l, err := net.Listen("tcp", this.frontend_addr)
		if err != nil {
			if this.OnErrFunc != nil {
				this.OnErrFunc(err)
			}
		} else {
			if this.OnListenSucc != nil {
				this.OnListenSucc()
			}
		}
		this.frontend = l
	}
	return this.frontend
}

/*
**

	连接remoteAddr, 完成后回调cb
*/
func (this *ProxyTcp) DoProxyConnect(frontend_conn net.Conn, remoteAddr string, cb func(frontend_conn, backend_conn net.Conn, err error)) {
	backend, err := net.Dial("tcp", remoteAddr)
	cb(frontend_conn, backend, err)
}

/*
**

	处理两个连接
*/
func (this *ProxyTcp) OnProxyConnCb(frontend_conn net.Conn, backend_conn net.Conn, err error) {
	if err != nil {
		frontend_conn.Close()
		if this.OnErrFunc != nil {
			this.OnErrFunc(err)
		}
		return
	}
	rec := &ProxyTcpRec{
		sessionid:     fmt.Sprintf("%s_%d", utils.RandKeyString(12), utils.NewObjectDNA()),
		Frontend_conn: frontend_conn,
		Backend_conn:  backend_conn,
	}
	this.Sessions.Set(rec.sessionid, rec)
	atomic.AddInt32(&this.OnlineCnt, 1)
	defer func() {
		this.onClose(rec)
	}()
	rec.FrontendAddr = frontend_conn.RemoteAddr().String()
	rec.BackendAddr = backend_conn.RemoteAddr().String()
	rec.StartT = time.Now()
	wait := new(sync.WaitGroup)
	wait.Add(2)

	if this.OnFrontendRecvBufConvert != nil {
		go DoIoCopyEx(backend_conn, frontend_conn, &rec.backend_recvsize, func(err error, tag int) {
			rec.RequestDisconnect(fmt.Sprintf("[代理前端[%s]->后端]:%s", frontend_conn.RemoteAddr().String(), err.Error()))
		}, func() {
			wait.Done()
		}, func() []byte {
			return make([]byte, 1024)
		}, func(buf []byte) {
			if this.OnFrontendRecv != nil {
				this.OnFrontendRecv(rec, buf)
			}
		}, func(buf []byte) []byte {
			if this.OnFrontendRecvBufConvert != nil {
				return this.OnFrontendRecvBufConvert(rec, buf)
			}
			return buf
		})
	} else {
		go DoIoCopy(backend_conn, frontend_conn, &rec.backend_recvsize, func(err error, tag int) {
			rec.RequestDisconnect(fmt.Sprintf("[代理前端[%s]->后端]:%s", frontend_conn.RemoteAddr().String(), err.Error()))
		}, func() {
			wait.Done()
		}, func() []byte {
			return make([]byte, 1024)
		}, func(buf []byte) {
			if this.OnFrontendRecv != nil {
				this.OnFrontendRecv(rec, buf)
			}
		})
	}

	go DoIoCopyEx(frontend_conn, backend_conn, &rec.frontend_sendsize, func(err error, tag int) {
		rec.RequestDisconnect(fmt.Sprintf("[代理后端[%s]->前端]:%s", backend_conn.RemoteAddr().String(), err.Error()))
	}, func() {
		wait.Done()
	}, func() []byte {
		return make([]byte, 1024)
	}, func(buf []byte) {
		if this.OnBackendRecv != nil {
			this.OnBackendRecv(rec, buf)
		}
	}, func(buf []byte) []byte {
		if this.OnBackendRecvBufConvert != nil {
			return this.OnBackendRecvBufConvert(rec, buf)
		}
		return buf
	})
	wait.Wait()
}

func (this *ProxyTcp) handleFrontendConn(frontend_conn net.Conn) {
	this.DoProxyConnect(frontend_conn, this.backend_addr, this.OnProxyConnCb)
}

func (this *ProxyTcp) onClose(rec *ProxyTcpRec) {
	atomic.AddInt32(&this.OnlineCnt, -1)
	this.Sessions.Remove(rec.sessionid)
	if this.OnDisconnected != nil {
		this.OnDisconnected(rec)
	}
}

func (this *ProxyTcp) GetBackendAddr() string {
	return this.backend_addr
}

func (this *ProxyTcp) GetFrontendAddr() string {
	return this.frontend_addr
}

func (this *ProxyTcp) Close() {
	this.terminated = true
	listener := this.frontend
	if listener != nil {
		listener.Close()
	}
}

func (this *ProxyTcp) GetActive() bool {
	return !this.terminated
}

func (this *ProxyTcp) BreakConnect() {
	listener := this.frontend
	if listener != nil {
		listener.Close()
	}
}

func (this *ProxyTcp) Open() {
	this.terminated = false
	this.closeW.Add(1)
	defer this.closeW.Done()
	for !this.terminated {
		listener := this.CheckFrontend()
		if listener == nil {
			time.Sleep(time.Second * 10)
			continue
		}
		client, err := listener.Accept()
		if err != nil {
			if this.OnErrFunc != nil {
				this.OnErrFunc(err)
			}
			this.frontend = nil
			continue
		}

		if this.OnHandleFrontendConnFunc != nil {
			if this.OnHandleFrontendConnFunc(client) == 0 {
				go this.handleFrontendConn(client)
			}
		} else {
			go this.handleFrontendConn(client)
		}
	}
}

func (this *ProxyTcp) SetFrontendAddr(val string) bool {
	if this.frontend_addr != val {
		this.frontend_addr = val
		return true
	}
	return false
}

func (this *ProxyTcp) SetBackendAddr(val string) bool {
	if this.backend_addr != val {
		this.backend_addr = val
		return true
	}
	return false
}
